import os
import subprocess
import argparse

def download_videos(base_url, cookies_file, output_prefix, start_p, end_p):
    """下载指定范围内的视频分P"""
    for p in range(start_p, end_p + 1):
        url = f"{base_url}&p={p}"
        output_name = f"{output_prefix}{p:03d}"
        
        print(f"\n正在下载分P{p}: {output_name}")
        print(f"URL: {url}")
        
        cmd = [
            "you-get",
            "--cookies=" + cookies_file,
            url,
            "-O", output_name
        ]
        
        try:
            subprocess.run(cmd, check=True)
            print(f"✓ 成功下载: {output_name}")
        except subprocess.CalledProcessError as e:
            print(f"✕ 下载失败: {output_name}, 错误: {e}")
            # 可以选择是否继续下载后续分P
            if input("是否继续下载其他分P? (y/n): ").lower() != 'y':
                break

def main():
    # 设置命令行参数
    parser = argparse.ArgumentParser(description='B站视频分P批量下载工具')
    parser.add_argument('base_url', help='视频基础URL')
    parser.add_argument('--cookies', default='cookies.txt', help='cookies文件路径 (默认: cookies.txt)')
    parser.add_argument('--prefix', default='视频_', help='输出文件名前缀 (默认: 视频_)')
    parser.add_argument('--start', type=int, default=1, help='开始的分P号 (默认: 1)')
    parser.add_argument('--end', type=int, required=True, help='结束的分P号 (必须指定)')
    parser.add_argument('--list', action='store_true', help='列出所有分P但不下载')
    
    args = parser.parse_args()

    # 检查cookies文件
    if not os.path.exists(args.cookies):
        print(f"错误: cookies文件 {args.cookies} 不存在!")
        return

    # 验证分P范围
    if args.start < 1:
        print("错误: 起始分P号必须大于0")
        return
    if args.end < args.start:
        print("错误: 结束分P号不能小于起始分P号")
        return

    if args.list:
        print(f"\n将要下载的分P列表 ({args.start} 到 {args.end}):")
        for p in range(args.start, args.end + 1):
            print(f"{args.prefix}{p:03d}")
        print(f"\n共 {args.end - args.start + 1} 个分P")
        return

    # 确认下载
    print(f"\n准备下载分P {args.start} 到 {args.end} 共 {args.end - args.start + 1} 个视频")
    print(f"输出文件名格式: {args.prefix}{args.start:03d} 到 {args.prefix}{args.end:03d}")
    confirm = input("是否开始下载? (y/n): ")
    if confirm.lower() != 'y':
        print("下载已取消")
        return

    # 开始下载
    download_videos(args.base_url, args.cookies, args.prefix, args.start, args.end)
    print("\n所有任务完成!")

if __name__ == "__main__":
    main()
    